import logging
from pathlib import Path
import torch
from typing import List
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from langchain_huggingface import HuggingFacePipeline
from tqdm import tqdm
from huggingface_hub import snapshot_download
import os

class ReviewSystem:
    def __init__(self):
        # 配置日志格式，添加中文提示
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s'
        )
        
        # 获取项目根目录
        self.root_dir = Path(__file__).parent.parent
        
        # 初始化路径
        self.markdown_dir = self.root_dir / "rag" / "markdown"
        self.model_path = Path(r"D:\python代码\AI大模型\参赛\Qwen2.5-3B-Instruct")
        
        # 检查并下载模型
        self.check_and_download_model()
        
        logging.info(f"正在加载模型: {self.model_path}")
        
        try:
            tokenizer = AutoTokenizer.from_pretrained(str(self.model_path), trust_remote_code=True)
            logging.info("分词器加载完成")
            
            model = AutoModelForCausalLM.from_pretrained(
                str(self.model_path),
                device_map="auto",
                trust_remote_code=True,
                torch_dtype=torch.bfloat16
            )
            logging.info("模型加载完成")
            
            # 创建pipeline
            pipe = pipeline(
                "text-generation",
                model=model,
                tokenizer=tokenizer,
                max_new_tokens=2048,
                temperature=0.7,
                top_p=0.8,
                repetition_penalty=1.05,
                device_map="auto",
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
                max_length=24576  # 设置一个合理的最大长度
            )
            logging.info("Pipeline创建完成")
            
            self.llm = HuggingFacePipeline(pipeline=pipe)
            logging.info("LangChain Pipeline初始化完成")
            
        except Exception as e:
            logging.error(f"模型加载失败: {str(e)}")
            raise

    def check_and_download_model(self):
        """检查并下载模型"""
        if not self.model_path.exists():
            logging.info(f"模型不存在，开始下载到: {self.model_path}")
            try:
                # 创建模型目录
                self.model_path.mkdir(parents=True, exist_ok=True)
                
                # 下载模型
                snapshot_download(
                    repo_id="Qwen/Qwen2.5-3B-Instruct",
                    local_dir=str(self.model_path),
                    local_dir_use_symlinks=False,
                    ignore_patterns=["*.md", "*.h5", "*.msgpack"]
                )
                logging.info("模型下载完成")
            except Exception as e:
                logging.error(f"模型下载失败: {str(e)}")
                raise
        else:
            logging.info("模型已存在，跳过下载")

    def load_markdown_content(self) -> str:
        """加载所有markdown文件内容"""
        all_content = []
        for md_file in tqdm(list(self.markdown_dir.rglob("*.md")), desc="加载文档"):
            try:
                content = md_file.read_text(encoding='utf-8')
                all_content.append(content)
            except Exception as e:
                logging.error(f"加载文件失败 {md_file}: {str(e)}")
        return "\n\n".join(all_content)

    def generate_review(self, prompt: str) -> str:
        """生成综述的核心方法"""
        try:
            # 加载所有markdown内容
            content = self.load_markdown_content()
            
            # 限制输入长度，取前20000个字符
            if len(content) > 20000:
                logging.warning(f"文档内容过长({len(content)}字符)，将截取前20000字符")
                content = content[:20000]
            
            # 构建完整的提示词
            full_prompt = f"""你是一个专业的学术助手。请基于以下文献内容，生成一个结构清晰、内容准确的综述。

文献内容：
{content}

要求：
{prompt}

请确保：
1. 内容客观准确
2. 结构清晰有序
3. 重点突出
4. 语言专业规范
5. 有理有据，适当引用原文关键内容

请开始生成综述：
"""
            # 生成回复
            response = self.llm.invoke(
                full_prompt,
                temperature=0.7,
                max_tokens=2048,
                top_p=0.8,
                repetition_penalty=1.05
            )
            return response
            
        except Exception as e:
            logging.error(f"生成综述失败: {str(e)}")
            return f"生成综述时发生错误: {str(e)}"

    def generate_concept_review(self, concept: str) -> str:
        """生成技术概念调研综述"""
        prompt = f"""
请对{concept}这个技术概念进行全面的调研综述。要求：
1. 概念定义与核心特征
2. 技术原理与关键组成
3. 应用场景与实际案例
4. 优势与局限性分析
5. 未来发展趋势

请保持客观准确，突出重点。
"""
        return self.generate_review(prompt)

    def generate_research_status(self, direction: str) -> str:
        """生成研究现状综述"""
        prompt = f"""
请对{direction}这个研究方向的现状进行系统综述。要求：
1. 研究背景与意义
2. 主要研究问题
3. 研究现状分析
4. 存在的问题与挑战
5. 未来研究方向

请保持客观中立的态度。
"""
        return self.generate_review(prompt)

    def generate_comparative_analysis(self, methods: List[str]) -> str:
        """生成多个方法的对比分析综述"""
        methods_str = "、".join(methods)
        prompt = f"""
请对{methods_str}这些方法进行系统的对比分析。要求：
1. 基本原理对比
2. 技术特点对比
3. 适用场景对比
4. 优劣势分析
5. 综合评估

请对比分析要客观公正。
"""
        return self.generate_review(prompt)

    def generate_technical_evolution(self, technology: str) -> str:
        """生成技术方法研究脉络综述"""
        prompt = f"""
请对{technology}的技术发展脉络进行系统梳理。要求：
1. 技术起源与背景
2. 关键发展阶段
3. 重要突破与创新
4. 技术演进路线
5. 未来发展趋势

请注重时间顺序和关键节点。
"""
        return self.generate_review(prompt) 